#!/bin/bash

function showhelp {
    echo 'usage : shs option dir_contain_sh or xxx.sh'
    echo 'e.g.  : shs -n 128 -E  -e node05 ABC_OPT'
    echo 'option: -[n total number of cores]'
    echo '        -[l node name. e.g. -l node01,node02] '
    echo '        -[e nodename to exclude. e.g. -e node01,node02]'
    echo '        -[q queue name to submit]'
    echo '        -[E exclusive mode]'
    echo '        -[p pbs or slurm]'
    echo '        -[s env set command]'
}
if [ $# -lt 1 ]; then
    showhelp
    exit 1
fi
while getopts ":g:n:l:e:s:q:E" opt
do
    case $opt in
        g) NGPUS=$OPTARG;;
        n) NPROC=$OPTARG;;
        l) NODE=$OPTARG;;
        e) EXCLUDE=$OPTARG;;
        q) QUEUE=$OPTARG;;
        s) ENVCOM=$OPTARG;;
        E) EXCLUSIVE=1;;
        ?) showhelp
        exit 1;;
    esac
done
shift $(($OPTIND - 1))


while [[ $# -ge 1 ]]
do
dos2unix $1
if [ -z "$NPROC" ]; then
    NPROC=64
fi

# check job manage system
PBS_FLAG=`which qsub 2> null`
SLURM_FLAG=`which sbatch 2> null`
if [ -n "$SLURM_FLAG" ];then
    JOB_NAME=${1%.*}
    NAME=sub_${JOBNAME}.sbatch
    INSUBDIR=''
    if [ -d "$1" ];then
        cd $1
        pwd
        INSUBDIR=1
        inpname=`ls *.sh`
    else
        inpname=$1
    fi
    echo '#!/bin/bash' > $NAME
    echo "#SBATCH -J $JOB_NAME" >> $NAME
    echo "#SBATCH -N 1" >> $NAME
    if [ -n "$NODE" ];then
        echo "#SBATCH -w $NODE" >> $NAME
    fi
    if [ -n "$EXCLUDE" ];then
        echo "#SBATCH --exclude $EXCLUDE" >> $NAME
    fi
    if [ -n "$EXCLUSIVE" ];then
        echo "#SBATCH --exclusive" >> $NAME
    fi
    if [ -n "$QUEUE" ];then
        echo "#SBATCH --partition=$QUEUE" >> $NAME
    fi
    echo "#SBATCH --ntasks-per-node=$NPROC" >> $NAME
    echo "#SBATCH --cpus-per-task=1" >> $NAME
    if [ -n "$NGPUS" ];then
        echo "#SBATCH -G $NGPUS" >> $NAME
    fi
    echo "source ~/.bashrc" >> $NAME
    if [ -n "$ENVCOM" ];then
        echo -e "$ENVCOM" >> $NAME
    fi
    echo "sh $inpname &> ${inpname}.out" >> $NAME
    sbatch $NAME
    if [ -n "$INSUBDIR" ];then
        cd ..
    fi
elif [ -n "$PBS_FLAG" ];then
    JOB_NAME=${1/.sh}
    NAME=sub_${1/.sh/}.pbs
    echo "#PBS -S /bin/bash" > $NAME
    if [ -n "$NODE" ];then
        echo "#PBS -l nodes=$NODE:ppn=$NPROC" >> $NAME
    else
        echo "#PBS -l nodes=1:ppn=$NPROC" >> $NAME
    fi
    echo "#PBS -N $JOB_NAME" >> $NAME
    echo "#PBS -j oe" >> $NAME
    echo 'cd $PBS_O_WORKDIR' >> $NAME
    echo "source ~/.bashrc" >> $NAME
    if [ -d "${1}" ]; then
        echo "cd $1" >> $NAME
        echo 'sh *.sh' >> $NAME
    else
        echo "sh $1 &> ${1%.*}.out" >> $NAME
    fi
    qsub $NAME
fi
shift 1
done 
